from datetime import datetime

import matplotlib.pyplot as plt
import torch
import torch.nn.functional as F
import torch.optim as optim
from torch_geometric.data import Data

from GNNModel import GNNModel

# Define bom_structure
# bom_structure = genBOM()
bom_structure = {'WGPart1': {
    'SBB3': ['SBB8', 'SBB4', 'SBB5'],
    'SBB4': ['SBB5', 'SBB4'],
    'SBB5': ['SBB5', 'SBB6'],
    'SBB6': ['SBB5', 'SBB3'],
    'SBB8': []
},
    'WGPart2': {},
    'WGPart3': {
        'SBB1': ['SBB2'],
        'SBB2': ['SBB1']
    },
    'WGPart4': {
        'SBB7': ['SBB7', 'SBB9']
    }
}

# Extract all components to components list
components = set(part for product, contents in bom_structure.items() for part in contents.keys())
components.update(
    subpart for product, contents in bom_structure.items() for subparts in contents.values() for subpart in
    subparts)
components = sorted(components)

# Create mapping from component name to index
component_to_index = {comp: index for index, comp in enumerate(components)}

# Build edge list
edge_list = []
for product, contents in bom_structure.items():
    for part, subparts in contents.items():
        for subpart in subparts:
            edge_list.append((component_to_index[part], component_to_index[subpart]))

# Convert edge list to PyTorch Geometric format
edge_index = torch.tensor(edge_list, dtype=torch.long).t().contiguous()

# Build feature matrix
n_components = len(components)
node_features = torch.eye(n_components)

# Build Data object
graph_data = Data(x=node_features, edge_index=edge_index)

# Initialize model
hidden_size = 8
out_size = n_components
model = GNNModel(node_features.size(1), hidden_size, out_size)


def train():
    # Train model
    optimizer = optim.Adam(model.parameters(), lr=0.01)
    model.train()
    losses = []
    for epoch in range(20):
        optimizer.zero_grad()
        out = model(graph_data.x, graph_data.edge_index)
        loss = F.mse_loss(out, torch.rand((n_components, out_size)))
        loss.backward()
        optimizer.step()
        losses.append(loss.item())
        if epoch % 10 == 0:
            print(f"Epoch {epoch}: Loss {loss.item()}")

    # 生成损失函数的变化图
    plt.plot(range(20), losses)
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.title('Loss over epochs')
    plt.savefig(f'loss_over_epochs_{datetime.now().strftime("%Y%m%d_%H%M%S")}.png', dpi=300, bbox_inches='tight')

    # Save model
    torch.save(model.state_dict(), "model.pth")


if __name__ == '__main__':
    train()
